# -*- coding: utf-8 -*-
# @Author  : Zhang.Jingyi

import os
import json
import numpy as np
import random

# ROOT_PATH = '/home/scsc02/jingyi/dataset/pedx'
ROOT_PATH = '/home/scsc02/jingyi/dataset/lidarcap'
label_2d_path = os.path.join(ROOT_PATH, 'labels/2d')
label_3d_pose_path = os.path.join(ROOT_PATH, 'labels/3d/pose')
label_3d_segment_path = os.path.join(ROOT_PATH, 'labels/3d/segment')
train_val_file = os.path.join(ROOT_PATH, 'train_val.json')
# train_txt = os.path.join(ROOT_PATH, 'train.txt')
# test_txt = os.path.join(ROOT_PATH,'test.txt')
if __name__ == '__main__':

    ##########################
    # generate train list
    ##########################
    train_human_list = {}
    test_human_list = {}
    test_human_list['00000000000000000000000000000011'] = {}
    test_human_list['146c364b4657488ea6575101b219803e'] = {}
    test_human_list['442f4341bf6945c9901729b7d14e8192'] = {}

    for date in ['20171130T2000', '20171207T2024']:
        data_dir = os.path.join(label_2d_path, date)
        for _, _, filenames in os.walk(data_dir):
            for filename in filenames:
                info = filename.split('_')
                smpl_para_json = os.path.join(
                    label_3d_pose_path, info[0], info[0] + '_' + info[2] + '_' + info[3])
                if not os.path.exists(smpl_para_json):
                    continue
                bbox_filepath = os.path.join(data_dir, filename)
                assert os.path.exists(bbox_filepath)
                mesh_file_path = os.path.join(
                    label_3d_pose_path, date, smpl_para_json.replace('.json', '.ply'))
                assert os.path.exists(mesh_file_path)
                point_file_path = os.path.join(
                    label_3d_segment_path, date, smpl_para_json.replace('.json', '.ply'))
                assert os.path.exists(point_file_path)
                pose_file_path = os.path.join(
                    label_3d_segment_path, date, smpl_para_json)
                assert os.path.exists(pose_file_path)
                with open(bbox_filepath, 'r') as f:
                    label_2d = json.load(f)
                tracking_id = label_2d['tracking_id']
                frame_id = label_2d['frame_id']
                if tracking_id in ['00000000000000000000000000000011',
                                   '146c364b4657488ea6575101b219803e',
                                   '442f4341bf6945c9901729b7d14e8192']:
                    test_human_list[tracking_id][frame_id] = {}
                    test_human_list[tracking_id][frame_id]['2D_bbox'] = bbox_filepath
                    test_human_list[tracking_id][frame_id]['3D_mesh'] = mesh_file_path
                    test_human_list[tracking_id][frame_id]['3D_points'] = point_file_path
                    test_human_list[tracking_id][frame_id]['3D_pose'] = pose_file_path

                else:
                    if tracking_id not in train_human_list.keys():
                        train_human_list[tracking_id] = {}
                    train_human_list[tracking_id][frame_id] = {}
                    train_human_list[tracking_id][frame_id]['2D_bbox'] = bbox_filepath
                    train_human_list[tracking_id][frame_id]['3D_mesh'] = mesh_file_path
                    train_human_list[tracking_id][frame_id]['3D_points'] = point_file_path
                    train_human_list[tracking_id][frame_id]['3D_pose'] = pose_file_path

    ##########################
    # generate validation list
    ##########################
    for date in ['20171130T2000', '20171207T2024']:
        test_smpl_path = os.path.join(label_3d_pose_path, date)
        for _, _, files in os.walk(test_smpl_path):
            for file in files:
                infos = file.split('_')
                if infos[-1].split('.')[-1] != 'ply':
                    frame = infos[1]
                    id = infos[2]
                    if os.path.exists(os.path.join(label_2d_path, date, date + '_ylw79D0_' + frame + '_' + id)):
                        continue
                    if os.path.exists(os.path.join(label_2d_path, date, date + '_blu79CF_' + frame + '_' + id)):
                        continue
                    if os.path.exists(os.path.join(label_2d_path, date, date + '_grn43E3_' + frame + '_' + id)):
                        continue
                    if os.path.exists(os.path.join(label_2d_path, date, date + '_red707B_' + frame + '_' + id)):
                        continue
                    mesh_file_path = os.path.join(
                        test_smpl_path, file.replace('.json', '.ply'))
                    assert os.path.exists(mesh_file_path)
                    pose_file_path = os.path.join(test_smpl_path, file)
                    assert os.path.exists(pose_file_path)
                    point_file_path = os.path.join(
                        label_3d_segment_path, date, file.replace('.json', '.ply'))
                    assert os.path.exists(point_file_path)
                    if id not in test_human_list.keys():
                        test_human_list[id] = {}
                    test_human_list[id][frame] = {}
                    test_human_list[id][frame]['2D_bbox'] = 'None'
                    test_human_list[id][frame]['3D_mesh'] = mesh_file_path
                    test_human_list[id][frame]['3D_points'] = point_file_path
                    test_human_list[id][frame]['3D_pose'] = pose_file_path

    train_val = dict(train_human_list=train_human_list,
                     test_human_list=test_human_list)
    with open(train_val_file, 'w') as f:
        json.dump(train_val, f)
